import tensorflow as tf
from stock_a.feature import feature_pb2
from stock_a.feature.abstract_feature import AbstractFeature
from stock_a.feature.stock_feature_source import StockFeatureSource
from stock_a.feature.global_feature_source import GlobalFeatureSource
from stock_a.common.db import stock_info
from stock_a.common.utils import collection_util

"""
近20日换手率序列，如果不足20日，以实际有的交易日数为准
"""
class SeqDayTurnoverFeature(AbstractFeature):
    def __init__(self, stock_feature_source: StockFeatureSource, global_feature_source: GlobalFeatureSource):
        super(SeqDayTurnoverFeature, self).__init__(stock_feature_source, global_feature_source)

    def feature_name(self):
        return "seq_day20_trunover"

    def feature_type(self):
        return tf.float32

    def feature_shape(self):
        return 20

    def set_pb_feature(self, idx: int, feature_pb: feature_pb2.feature):
        start_idx= max(idx - 19, 0)
        for i in range(start_idx, idx + 1):
            feature_pb.seq_day20_trunover.append(self.stock_feature_source.day_k_info_list[i][stock_info.TURNOVER_INDEX])

    def get_pb_feature(self, feature_pb: feature_pb2.feature):
        seq_day20_trunover = feature_pb.seq_day20_trunover
        # 如果长度不足20，那么用0 padding
        if seq_day20_trunover is None:
            seq_day20_trunover = [0.0 for _ in range(20)]
        elif len(seq_day20_trunover) < 20:
            seq_day20_trunover = collection_util.pad_list_front(list(seq_day20_trunover), 20)
        return seq_day20_trunover


"""
近20日涨跌幅序列，如果不足20日，以实际有的交易日数为准
"""
class SeqDayFluctuateFeature(AbstractFeature):
    def __init__(self, stock_feature_source: StockFeatureSource, global_feature_source: GlobalFeatureSource):
        super(SeqDayFluctuateFeature, self).__init__(stock_feature_source, global_feature_source)

    def feature_name(self):
        return "seq_day20_fluctuate"

    def feature_type(self):
        return tf.float32

    def feature_shape(self):
        return 20

    def set_pb_feature(self, idx: int, feature_pb: feature_pb2.feature):
        start_idx = max(idx - 19, 0)
        for i in range(start_idx, idx + 1):
            feature_pb.seq_day20_fluctuate.append(
                self.stock_feature_source.day_k_info_list[i][stock_info.FLUCTUATE_HFQ_BY_FACTOR_INDEX])

    def get_pb_feature(self, feature_pb: feature_pb2.feature):
        seq_day20_fluctuate = feature_pb.seq_day20_fluctuate
        # 如果长度不足20，那么用0 padding
        if seq_day20_fluctuate is None:
            seq_day20_fluctuate = [0.0 for _ in range(20)]
        elif len(seq_day20_fluctuate) < 20:
            seq_day20_fluctuate = collection_util.pad_list_front(list(seq_day20_fluctuate), 20)
        return seq_day20_fluctuate

"""
近20日振幅序列，如果不足20日，以实际有的交易日数为准
"""
class SeqDayAmplitudeFeature(AbstractFeature):
    def __init__(self, stock_feature_source: StockFeatureSource, global_feature_source: GlobalFeatureSource):
        super(SeqDayAmplitudeFeature, self).__init__(stock_feature_source, global_feature_source)

    def feature_name(self):
        return "seq_day20_amplitude"

    def feature_type(self):
        return tf.float32

    def feature_shape(self):
        return 20

    def set_pb_feature(self, idx: int, feature_pb: feature_pb2.feature):
        start_idx = max(idx - 19, 0)
        for i in range(start_idx, idx + 1):
            feature_pb.seq_day20_amplitude.append(
                self.stock_feature_source.day_k_info_list[i][stock_info.AMPLITUDE_HFQ_BY_FACTOR_INDEX])

    def get_pb_feature(self, feature_pb: feature_pb2.feature):
        seq_day20_amplitude = feature_pb.seq_day20_amplitude
        # 如果长度不足20，那么用0 padding
        if seq_day20_amplitude is None:
            seq_day20_amplitude = [0.0 for _ in range(20)]
        elif len(seq_day20_amplitude) < 20:
            seq_day20_amplitude = collection_util.pad_list_front(list(seq_day20_amplitude), 20)
        return seq_day20_amplitude
